from flask import Flask, request, jsonify
from userVanna import userVanna
import uuid
import logging
import threading
import json
import os


app = Flask(__name__)
app.secret_key = '1345456'  # 设置一个秘密密钥用于会话加密

# 配置日志记录
logging.basicConfig(level=logging.DEBUG)

# 存储实例的字典
instances = {}

# 存储预训练日志的字典
pre_training_logs = {}

# 预训练信息文件路径
PRE_TRAINED_FILE = 'pre_trained.json'
#预训练sql文件路径
PRE_TRAINED_SQL_FILE = 'pre_trained_sql.json'

def load_pre_trained_info(file_name):
    if os.path.exists(file_name):
        with open(file_name, 'r') as f:
            return json.load(f)
    return {}

def save_pre_trained_info(file_name,info):
    with open(file_name, 'w') as f:
        json.dump(info, f)

@app.route('/connect', methods=['POST'])
def connect():
    data = request.json
    sql_name = data['sql_name']
    ip = request.remote_addr
    pre_trained_info = load_pre_trained_info(PRE_TRAINED_FILE)
    #获取训练模式 1-训练ddl；2-训练sql；3-训练文档；4-问答模式
    role = data['role']

    if role == '1':
        # 检查是否已经有预训练好的实例
        if ip in pre_trained_info and sql_name in pre_trained_info[ip]:
            dbname = data.get('dbname', '')
            if dbname in pre_trained_info[ip][sql_name] and pre_trained_info[ip][sql_name][dbname]:
                # 使用预训练好的实例
                instance = userVanna(sql_name, user_id=data['host'])
                instance.connect(**data)
                instance_id = str(uuid.uuid4())
                instances[instance_id] = instance
                return   instance_id

        # 如果没有预训练好的实例，则创建新的实例并进行预训练
        if sql_name == 'mysql':
            host = data['host']
            dbname = data['dbname']
            user = data['user']
            password = data['password']
            port = data['port']
            # role = data['role']
            instance = userVanna(sql_name, user_id=data['host'])
            instance.connect(host=host, dbname=dbname, user=user, password=password, port=port)
        elif sql_name == 'sqlite':
            host = data['host']
            port = data['port']
            dbname = data['dbname']
            # role = data['role']
            instance = userVanna(sql_name, user_id=data['host'])
            instance.connect(host=host, port=port, dbname=dbname)
        elif sql_name == 'snowflake':
            host = data['host']
            user = data['user']
            password = data['password']
            dbname = data['dbname']
            # role = data['role']
            instance = userVanna(sql_name, user_id=data['host'])
            instance.connect(host=host, user=user, password=password, dbname=dbname, role=role)

        instance_id = str(uuid.uuid4())  # 使用UUID作为实例的唯一标识
        instances[instance_id] = instance

        #检查预训练状态
        if not pre_trained_info.get(ip, {}).get(sql_name, {}).get(dbname, False):
            pre_training_logs[ip] = []
            # 在单独的线程中启动预训练
            threading.Thread(target=pre_train_async, args=(instance, ip, sql_name, dbname)).start()
        # # 检查预训练状态
        # if not pre_trained_info.get(ip, {}).get(sql_name, False):
        #     pre_training_logs[ip] = []
        #     # 在单独的线程中启动预训练
        #     threading.Thread(target=pre_trainsql_async, args=(instance, ip, sql_name, dbname)).start()
        return instance_id
    elif role == '2':
        #创建实例
        if sql_name == 'mysql':
            host = data['host']
            dbname = data['dbname']
            user = data['user']
            password = data['password']
            port = data['port']
            # role = data['role']
            instance = userVanna(sql_name, user_id=data['host'])
            instance.connect(host=host, dbname=dbname, user=user, password=password, port=port)

        instance_id = str(uuid.uuid4())  # 使用UUID作为实例的唯一标识
        instances[instance_id] = instance

        #检查预训练状态
        if  pre_trained_info.get(ip) !=None:
            pre_training_logs[ip] = []
            # 在单独的线程中启动预训练
            threading.Thread(target=pre_trainsql_async, args=(instance, ip, sql_name, dbname)).start()


        return instance_id
    elif role == '4':
        #创建实例
        if sql_name == 'mysql':
            host = data['host']
            dbname = data['dbname']
            user = data['user']
            password = data['password']
            port = data['port']
            # role = data['role']
            instance = userVanna(sql_name, user_id=data['host'])
            instance.connect(host=host, dbname=dbname, user=user, password=password, port=port)

        instance_id = str(uuid.uuid4())  # 使用UUID作为实例的唯一标识
        instances[instance_id] = instance

        #检查预训练状态
        if  pre_trained_info.get(ip) !=None:
            pre_training_logs[ip] = []
            # 在单独的线程中启动预训练
            threading.Thread(target=pre_trainask_async, args=(instance, ip, sql_name, dbname)).start()


        return instance_id           


def pre_train_async(instance, ip, sql_name, db_name):
    try:
        instance.pre_train(log_callback=lambda log: pre_training_logs[ip].append(log))
        pre_trained_info = load_pre_trained_info(PRE_TRAINED_FILE)
        if ip not in pre_trained_info:
            pre_trained_info[ip] = {}
        if sql_name not in pre_trained_info[ip]:
            pre_trained_info[ip][sql_name] = {}
        pre_trained_info[ip][sql_name][db_name] = True
        save_pre_trained_info(pre_trained_info)
    except Exception as e:
        logging.error(f"预训练过程中发生错误: {e}")

def pre_trainsql_async(instance, ip, sql_name, db_name):
    try:
        pre_trained_info = load_pre_trained_info(PRE_TRAINED_SQL_FILE)
        train_sql = []
        if ip in pre_trained_info:
            sql=pre_trained_info[ip]['sql'] 
            #将sql字符串安装分号分割成列表装到train_sql中
            train_sql = sql.split(';')
        instance.sql_train(train_sql)
        
    except Exception as e:
        logging.error(f"预训练过程中发生错误: {e}")

def pre_trainask_async(instance, ip, sql_name, db_name):
    try:
        instance.sql_question_train('sql_question_train.xlsx')       
    except Exception as e:
        logging.error(f"预训练过程中发生错误: {e}")

@app.route('/get_sql_code', methods=['POST'])
def get_sql_code():
    data = request.json
    question = data['question']
    instance_id = data['instance_id']
    if instance_id:
        instance = instances.get(instance_id)
        if instance:
            return instance.generate_sql_code(question)
    return jsonify({'error': 'Instance not found'}), 404

@app.route('/get_answer', methods=['POST'])
def get_answer():
    data = request.json
    question = data['question']
    instance_id = data['instance_id']
    if instance_id:
        instance = instances.get(instance_id)
        if instance:
            return instance.ask(question)
    return jsonify({'error': 'Instance not found'}), 404

if __name__ == '__main__':
    app.run(use_reloader=False, threaded=True, debug=True, host='0.0.0.0', port=5600)

